# 导入rnn_utils模块，用于处理可变长度序列的填充和排序
import os

import torch.nn.utils.rnn as rnn_utils
# 导入Dataset和DataLoader模块，用于加载和处理数据集
from torch.utils.data import Dataset, DataLoader
import pickle  # 导入pickle模块，用于序列化和反序列化Python对象

from .dataset import ChatDataset  # 导入自定义的数据集类 

# 这里还是说明一下为啥(from dataset----->from .dataset)，在python中，不同目录位置的代码导包容易出现问题，
# 现在我改为了from .dataset 那这个代码就不能单独运行了。
# ✅ 适用于从 根目录的train.py 调用时运行良好
# ❌ 不适用于右键运行 data_loader.py  你也可以看到这个最下面的 if __name__ == '__main__'测试代码我删了
# （会报错 ImportError: attempted relative import with no known parent package）

def collate_fn(batch):
    """
    自定义的collate_fn函数，用于将数据集中的样本进行批处理
    :param batch: 样本列表
    :return: 经过填充的输入序列张量和标签序列张量
    """
    # 对输入序列进行填充，使其长度一致
    input_ids = rnn_utils.pad_sequence(batch, batch_first=True, padding_value=0)
    
    # 对标签序列进行填充，使其长度一致
    labels = rnn_utils.pad_sequence(batch, batch_first=True, padding_value=-100)

    return input_ids, labels




def load_dataset(train_path, valid_path):
    # print('进入函数')
    """
    加载训练集和验证集
    :param train_path: 训练数据集路径
    :return: 训练数据集和验证数据集
    """
    with open(train_path, "rb") as f:
        train_input_list = pickle.load(f)  # 从文件中加载输入列表

    with open(valid_path, "rb") as f:
        valid_input_list = pickle.load(f)  # 从文件中加载输入列表

    train_dataset = ChatDataset(train_input_list, 300)  # 创建训练数据集对象
    val_dataset = ChatDataset(valid_input_list, 300)  # 创建验证数据集对象
    return train_dataset, val_dataset  # 返回训练数据集和验证数据集

def get_dataloader(train_path, valid_path):
    train_dataset, valid_dataset = load_dataset(train_path, valid_path)
    # print(f'len-->{len(train_dataset)}')
    # print(f'len-->{len(valid_dataset)}')

    # 创建训练数据集的DataLoader对象
    train_dataloader = DataLoader(train_dataset,
                                  batch_size=4,
                                  shuffle=True,
                                  collate_fn=collate_fn,
                                  drop_last=True)

    # 创建验证数据集的DataLoader对象
    valid_dataloader = DataLoader(valid_dataset,
                                  batch_size=4,
                                  shuffle=True,
                                  collate_fn=collate_fn,
                                  drop_last=True)

    return train_dataloader, valid_dataloader
    

def load_chat_dataset(chat_path):
    # print('进入函数')
    """
    加载训练集和验证集
    :param train_path: 训练数据集路径
    :return: 训练数据集和验证数据集
    """
    with open(chat_path, "rb") as f:
        chat_input_list = pickle.load(f)  # 从文件中加载输入列表

    chat_dataset = ChatDataset(chat_input_list, 100)  # 创建训练数据集对象
    return chat_dataset  # 返回训练数据集


def get_chat_dataloader(chat_path):
    chat_dataset = load_chat_dataset(chat_path)

    # 创建训练数据集的DataLoader对象
    chat_dataloader = DataLoader(chat_dataset,
                                  batch_size=4,
                                  shuffle=True,
                                  collate_fn=collate_fn,
                                  drop_last=False)

    return chat_dataloader

